import numpy as np


def box_sorted(boxes_in, image=[], default_iou=0.5, debug_show=False):
    """
    将box按照行（行内从左到右）从上到下排序
    :param boxes_in: list, [[x,y,w,h],[x,y,w,h],..], 该图每一个文本框的位置,x、y为文本框左上角点坐标,w、h为文本框的宽高
    :param image: ndarray, uint8, bgr
    :param debug_show: bool, True:开启结果展示, False:关闭结果展示（默认）
    :param default_iou: ndarray, float32:需要考虑行之间没有严格的界限情况（即错位情况）：通过列上重叠度来处理
    :return:
    boxes_out: 格式同输入boxes_in
    boxes_out_each_line, list, [[], [],...],列表中每个[]数据格式同输入boxes_in, 表示同一行的boxes,列表顺序对应文本行顺序
    """

    if not boxes_in:
        # 保护处理：输入为空的情况
        return boxes_in, boxes_in

    # 根据box中心行坐标先从上到下(y:box[1], h:box[3])排序
    boxes = sorted(boxes_in, key=(lambda box: [box[1] + box[3] / 2, box[0]]), reverse=False)

    # 统计平均行高
    heights = [box[3] for box in boxes]
    box_h_mean = np.mean(heights)

    # 求每个box的中心高度
    centers_y = [box[1] + box[3] / 2 for box in boxes]

    # 根据平均行高(决定了允许的倾斜程度,但不可设置太大，容易错行)归并box到各自的行
    centers_y_ls = []
    idxs_ls = []
    for idx in range(len(centers_y)):
        center = centers_y[idx]
        if not centers_y_ls:
            centers_y_ls.append([center])
            idxs_ls.append([idx])
            continue

        if abs(center - centers_y_ls[-1][-1]) < box_h_mean:
            over_lap_state = False

            # 需要考虑行之间没有严格的界限情况（即错位情况）：通过列上重叠度来处理
            box_c = boxes[idx]
            for box_idx in idxs_ls[-1]:
                box = boxes[box_idx]
                x_left = max(box_c[0], box[0])
                x_right = min(box_c[0] + box_c[2], box[0] + box[2])
                w = max(min(box_c[2], box[2]), 1)
                iou = (x_right - x_left) / w
                if iou > default_iou:
                    over_lap_state = True
                    break

            if over_lap_state:
                # 列上有重叠，需另起一行
                centers_y_ls.append([center])
                idxs_ls.append([idx])
            else:
                centers_y_ls[-1].append(center)
                idxs_ls[-1].append(idx)
        else:
            centers_y_ls.append([center])
            idxs_ls.append([idx])

    boxes_out = []
    boxes_out_each_line = []
    for idxs in idxs_ls:
        boxes_this = boxes[idxs[0]:idxs[-1] + 1]
        # 行内从左到右(以框的中心，稳定性相对更高)排序(此步骤之前经过开始的排序处理，box的顺序已经是从上到下了)
        boxes_this = sorted(boxes_this, key=(lambda box: [box[0] + box[2] / 2]), reverse=False)

        # 行内排完序后还需要重新判断下前后box上下行高关系(一行存在多个box情况)：处理中间数据被分到另一行
        if len(boxes_this) == 1:
            boxes_mult = [boxes_this]
        else:
            boxes_mult = [[boxes_this[0]]]
            for i in range(1, len(boxes_this)):
                box_c = boxes_this[i]
                c_used_state = False
                for line_idx, boxes_line in enumerate(boxes_mult):
                    # 因为行内排了序，只与已分开行的最后一个box框比较
                    box_b = boxes_line[-1]

                    y_center_c = box_c[1] + box_c[3] / 2
                    y_center_b = box_b[1] + box_b[3] / 2

                    if abs(y_center_c - y_center_b) < box_h_mean:
                        boxes_mult[line_idx].append(box_c)
                        c_used_state = True
                        break

                # 若行内当前box未找到可归并的行，则另起一行
                if not c_used_state:
                    boxes_mult.append([box_c])

            # 若行被重新拆分成多行，需要行重新排序
            if len(boxes_mult) > 1:
                # 根据行内所有box中心行位置均值
                boxes_mult = sorted(boxes_mult, key=(lambda boxes: [np.mean([box[1] + box[3] / 2 for box in boxes])]),
                                    reverse=False)

        for boxes_single_line in boxes_mult:
            boxes_out += boxes_single_line
            boxes_out_each_line.append(boxes_single_line)

            if debug_show:
                # 查看行归并是否正确
                import cv2
                img_show = image.copy()
                for pos in boxes_single_line:
                    pt1 = (pos[0], pos[1])
                    pt2 = (pos[0] + pos[2], pos[1] + pos[3])
                    cv2.rectangle(img_show, pt1, pt2, (255, 0, 0), 2)
                if img_show.shape[0] > 900 or img_show.shape[1] > 1200:
                    cv2.namedWindow('img_show_out', 0)
                    cv2.resizeWindow("img_show_out", 1200, 900)
                cv2.imshow("img_show_out", img_show)
                cv2.waitKey()

    return boxes_out, boxes_out_each_line


def box_sorted2(boxes_in, image, debug_show=False):
    """
    将box按照行（行内从左到右）从上到下排序
    :param boxes_in: list, [[x,y,w,h],[x,y,w,h],..], 该图每一个文本框的位置,x、y为文本框左上角点坐标,w、h为文本框的宽高
    :param image: ndarray, uint8, bgr
    :param debug_show: bool, True:开启结果展示, False:关闭结果展示（默认）
    :return:
    boxes_out: 根据行输出， 每行结果对应一个box框，整体格式同输入boxes_in
    boxes_each_line_less: 每行少于固定个数的box集中在一起，格式同输入boxes_in
    """

    if not boxes_in:
        # 保护处理：输入为空的情况
        return [], []

    # 根据box中心行坐标先从上到下(y:box[1], h:box[3])排序
    boxes = sorted(boxes_in, key=(lambda box: [box[1] + box[3] / 2, box[0]]), reverse=False)

    # 统计平均行高
    heights = [box[3] for box in boxes]
    box_h_mean = np.mean(heights)

    # 求每个box的中心高度
    centers_y = [box[1] + box[3] / 2 for box in boxes]

    # 根据平均行高(决定了允许的倾斜程度,但不可设置太大，容易错行)归并box到各自的行
    centers_y_ls = []
    idxs_ls = []
    for idx in range(len(centers_y)):
        center = centers_y[idx]
        if not centers_y_ls:
            centers_y_ls.append([center])
            idxs_ls.append([idx])
            continue

        if abs(center - centers_y_ls[-1][-1]) < box_h_mean:
            centers_y_ls[-1].append(center)
            idxs_ls[-1].append(idx)
        else:
            centers_y_ls.append([center])
            idxs_ls.append([idx])

    boxes_out = []
    boxes_each_line_less = []
    less_n = 3
    for idxs in idxs_ls:
        boxes_this = boxes[idxs[0]:idxs[-1] + 1]

        # 将忽略掉少于less_n个数的hbox行,用于后续改变这些box属性（如header转key）
        if len(boxes_this) < less_n:
            boxes_each_line_less += boxes_this
            continue

        # 行内从左到右(以框的中心，稳定性相对更高)排序(此步骤之前经过开始的排序处理，box的顺序已经是从上到下了)
        boxes_this = sorted(boxes_this, key=(lambda box: [box[0] + box[2] / 2]), reverse=False)

        x_min, y_min = np.min(boxes_this, axis=0)[:2]

        x2_y2_ls = [[box[0] + box[2], box[1] + box[3]] for box in boxes_this]
        x_max, y_max = np.max(x2_y2_ls, axis=0)[:2]

        # 行默认贯通整个图像的宽度
        x_min = 0
        x_max = image.shape[1]

        w = x_max - x_min
        h = y_max - y_min

        boxes_out.append([int(x_min), int(y_min), int(w), int(h)])

        if debug_show:
            import cv2
            img_show_out = image.copy()
            for box in boxes:
                pt1 = (box[0], box[1])
                pt2 = (box[0] + box[2], box[1] + box[3])
                cv2.rectangle(img_show_out, pt1, pt2, (0, 255, 0), 2)

            cv2.imshow("boxes_header", img_show_out)
            cv2.waitKey()

            # 查看行归并是否正确
            img_show_out = image.copy()
            pt1 = (int(x_min), int(y_min))
            pt2 = (int(x_max), int(y_max))
            cv2.rectangle(img_show_out, pt1, pt2, (255, 0, 0), 2)
            if img_show_out.shape[0] > 900 or img_show_out.shape[1] > 1200:
                cv2.namedWindow('img_show_out', 0)
                cv2.resizeWindow("img_show_out", 1200, 900)
            cv2.imshow("img_show_out", img_show_out)
            cv2.waitKey()

    return boxes_out, boxes_each_line_less


if __name__ == '__main__':
    # demo
    import json
    import cv2

    img_path = "/kv-data/sw/intelliphile2/7bbad868-9110-11ea-8944-a85e455eb6a1-perspective.png"
    boxes_in_path = "/kv-data/alg/graw/致泛新版本/new_test_datas2/记录/记录6_word_boxes.json"
    with open(boxes_in_path, "r") as f:
        boxes_in = json.load(f)
    image = cv2.imread(img_path)

    box_sorted(boxes_in, image, debug_show=True)

    print("over")